
from keras.layers.merge import add
from keras.layers.normalization import BatchNormalization
from keras.regularizers import l2
from keras import backend as K
from keras.layers import Input,Activation,Flatten,Dense,Conv2D,MaxPool2D,Conv2DTranspose,concatenate
from keras.models import Model
from keras.losses import binary_crossentropy
from keras.callbacks import ModelCheckpoint,EarlyStopping,ReduceLROnPlateau,TensorBoard

class RSSegVGGNet(object):
    @staticmethod
    def build():
        # define model
        image_input = Input(shape=(288,288,3), name='image_input') # (None, 288,288,3)

        # blobk 1
        conv1 = Conv2D(64, 3, activation='relu', padding='same')(image_input)
        conv1 = Conv2D(64, 3, activation='relu', padding='same')(conv1) # (None, 288, 288, 64)
        conv1pool = MaxPool2D(2,2)(conv1) # (None, 144,144,64)

        # block 2
        conv2 = Conv2D(128, 3, activation='relu', padding='same')(conv1pool)
        conv2 = Conv2D(128, 3, activation='relu', padding='same')(conv2) # (None, 144,144,128)
        conv2pool = MaxPool2D(2,2)(conv2) # (None, 72,72,128)

        # block 3
        conv3 = Conv2D(256, 3, activation='relu', padding='same')(conv2pool)
        conv3 = Conv2D(256, 3, activation='relu', padding='same')(conv3)
        conv3 = Conv2D(256, 3, activation='relu', padding='same')(conv3) # (None, 72,72,256)
        conv3pool = MaxPool2D(2,2)(conv3)# (None, 36,36,256)

        # block 4
        conv4 = Conv2D(512, 3, activation='relu', padding='same')(conv3pool)
        conv4 = Conv2D(512, 3, activation='relu', padding='same')(conv4)
        conv4 = Conv2D(512, 3, activation='relu', padding='same')(conv4) # (None, 36,36,512)
        conv4pool = MaxPool2D(2,2)(conv4) # (None,18,18,512)

        # block 5
        conv5 = Conv2D(512, 3, activation='relu', padding='same')(conv4pool)
        conv5 = Conv2D(512, 3, activation='relu', padding='same')(conv5)
        conv5 = Conv2D(512, 3, activation='relu', padding='same')(conv5)  # (None, 18, 18, 512)
        conv5pool = MaxPool2D(2,2)(conv5) # (None, 9,9,512)

        # center
        center = Conv2D(512, 3, activation='relu', padding='same')(conv5pool) # (None, 9,9,256)

        # decoder
        db5 = Conv2DTranspose(256, 3, strides=(2,2), activation='relu', padding='same')(center) #(None, 18, 18, 256)
        db5 = concatenate([db5, conv5], axis=3)
        db5 = Conv2D(512, 3, activation='relu', padding='same')(db5)#(None, 18, 18, 512)

        db4 = Conv2DTranspose(256, 3, strides=(2,2), activation='relu', padding='same')(db5) #(None, 36, 36, 256)
        db4 = concatenate([db4, conv4], axis=3)
        db4 = Conv2D(512, 3, activation='relu', padding='same')(db4) #(None, 36, 36, 512)

        db3 = Conv2DTranspose(128, 3, strides=(2,2), activation='relu', padding='same')(db4) # (None, 72,72,128)
        db3 = concatenate([db3, conv3], axis=3)
        db3 = Conv2D(256, 3, activation='relu', padding='same')(db3) #(None, 72, 72, 256)

        db2 = Conv2DTranspose(64,  3, strides=(2,2), activation='relu', padding='same')(db3) # (None, 144,144,64)
        db2 = concatenate([db2, conv2], axis=3)
        db2 = Conv2D(128, 3, activation='relu', padding='same')(db2) #(None, 144, 144, 128)

        db1 = Conv2DTranspose(32,  3, strides=(2,2), activation='relu', padding='same')(db2) # (None, 288,288,32)
        db1 = concatenate([db1,conv1],axis = 3)
        db1 = Conv2D(1,  3, activation='sigmoid', padding='same')(db1) # (None, 288, 288, 1)

        model = Model(inputs=image_input, outputs=db1)
        return model



